package com.rpp.study.filter;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.http.HttpStatus;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Mono;

import javax.annotation.PostConstruct;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicLong;

/**
 * @Description
 * @Author rpp
 * @Date 2020/8/30 8:58 下午
 */
@Slf4j
@Component
public class IpFilter implements GlobalFilter, Ordered {

    /**
     * 限流工具，限制单个客户端IP在最近X分钟内请求注册接口不能超过Y次
     */
//    public static ConcurrentHashMap<String, RateLimiter> resourceRateLimiter = new ConcurrentHashMap();

    /**
     * 请求时间间隔
     */
    @Value("${user.request.timePeriod}")
    private Integer timePeriod;

    /**
     * 最大请求次数
     */
    @Value("${user.request.maxTimes}")
    private Integer maxTimes;

    /**
     * 限流工具，限制单个客户端IP在最近X分钟内请求注册接口不能超过Y次
     */
    private LoadingCache<String, AtomicLong> counter;

    @PostConstruct
    public void init() {
        counter = CacheBuilder.newBuilder()
                .expireAfterWrite(timePeriod, TimeUnit.MINUTES)
                .build(new CacheLoader<>() {
                    @Override
                    public AtomicLong load(String key) {
                        return new AtomicLong(0);
                    }
                });
    }

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        ServerHttpRequest request = exchange.getRequest();
        ServerHttpResponse response = exchange.getResponse();

        String clientIp = request.getRemoteAddress().getHostString();
        String path = request.getURI().getPath();

        String key = clientIp + ":" + path;
        log.info("请求ip:{}, path:{}, timePeriod:{}, maxTimes:{}", clientIp, path, timePeriod, maxTimes);

        Boolean result = true;
        try {
            if (counter.get(key).incrementAndGet() > maxTimes) {
                result = false;
            }
        } catch (ExecutionException e) {
            e.printStackTrace();
        }
        //初始化限流工具，限制单个客户端IP在最近X分钟内请求注册接口不能超过Y次
//        resourceRateLimiter.putIfAbsent(key, RateLimiter.create(maxTimes, timePeriod, TimeUnit.MINUTES));
//        RateLimiter rateLimiter = resourceRateLimiter.get(key);
//        boolean result = rateLimiter.tryAcquire();

        if (!result) {
            // 拒绝访问
            response.setStatusCode(HttpStatus.SEE_OTHER);
            log.info("{} 时间内，ip:{}, path:{},请求超过最大次数：{}", timePeriod, clientIp, path, maxTimes);
            String data = "您频繁进⾏注册，请求已被拒绝";
            DataBuffer wrap = response.bufferFactory().wrap(data.getBytes());
            response.getHeaders().add("Content-Type", "text/plain;charset=UTF-8");
            return response.writeWith(Mono.just(wrap));
        }

        // 合法请求，放行，执行后续的过滤器
        return chain.filter(exchange);
    }

    @Override
    public int getOrder() {
        return 1;
    }
}
